import base64
import time
import os

import requests
from PIL import Image
import numpy as np

import model_precit
import util.img_utils as img_utils
import model_type.qunar.my_config as my_config


def removeFrame(img, width):
    '''
    :param img:
    :param width: 边框的宽度
    :return:
    '''
    w, h = img.size
    pixdata = img.load()
    for x in range(width):
        for y in range(0, h):
            pixdata[x, y] = 255
    for x in range(w - width, w):
        for y in range(0, h):
            pixdata[x, y] = 255
    for x in range(0, w):
        for y in range(0, width):
            pixdata[x, y] = 255
    for x in range(0, w):
        for y in range(h - width, h):
            pixdata[x, y] = 255
    return img


def cutWhiteArea(im, threshold=187, num=4):
    width, height = im.size
    ima = np.asarray(im)
    minx = height
    miny = width
    maxx = 0
    maxy = 0
    for x in range(0, height):
        for y in range(0, width):
            if ima[x, y] < threshold:
                if minx > x:
                    minx = x
                if miny > y:
                    miny = y
                if maxx < x:
                    maxx = x
                if maxy < y:
                    maxy = y
    print("h:{} ,w:{}".format(maxx - minx, maxy - miny))
    new_img = ima[minx:maxx + 1, miny:maxy + 1]
    width = new_img.shape[1]
    a = [0]
    for i in range(1, num + 1):
        a.append(int(width / num) * i)
    # 前后增加5个像素
    # a[0] = a[0] + 5
    # a[-1] = a[-1] - 5

    nim = []
    for i in range(0, num):
        w1, w2 = getlr(a, width, i)
        nim.append(Image.fromarray(new_img[:, w1:w2].astype('uint8')).convert('L')
                   .resize(my_config.img_size, resample=Image.ANTIALIAS))

    return nim


def getlr(a, width, i):
    # add_p = my_config.add_p
    # k = add_p * 2
    # k1 = add_p
    # if a[i+1]+add_p > width:
    #     k1 = width - a[i+1]
    #
    # if a[i] <= k1:
    #     w1 = 0
    # else:
    #     w1 = a[i] - k1
    # k = k - (a[i] - w1)
    # if a[+1] + k >= width:
    #     w2 = width
    # else:
    #     w2 = a[i + 1] + k
    # return w1, w2
    return a[i], a[i + 1]


def formatName(name):
    '''
    名称处理
    :param name:
    :return:
    '''
    if name:
        if '-' in name:
            code = name.split('-')[0]
            name = code.lower()
        else:
            name = name.split('.')[0].lower()
    return list(name)


def cut_img(im):
    # 基本处理，灰度处理，提升识别准确率
    im = im.convert("L")
    # 去掉边框
    im = removeFrame(im, 1)
    # 切去白色区域并且分成四个字母
    nim = cutWhiteArea(im)
    return nim

def format_img(im, name, newFilepath='label/tmp'):
    '''
    格式化图片
    '''
    oldName = name.split('-')[0].split('.')[0].lower()

    if not os.path.exists(newFilepath):
        os.makedirs(newFilepath)
    nim = cut_img(im)
    # 格式化图片名
    name = formatName(name)
    if len(nim) != len(name):
        name = ['z'] * len(nim)
    for i in range(len(nim)):

        nim[i].save(newFilepath + '/' + name[i] + '-' +
                    oldName + '_' + (str(time.time())).replace(".", "") + ".png", 'png')


def format_img_to_model(im):
    '''
    为模型格式化参数
    :param im:
    :return:
    '''
    captcha_array = np.array(im).astype('float32') / 255  # 向量化
    width, height = my_config.img_size
    captcha_array = captcha_array.reshape((width, height, -1))
    return captcha_array


def gen_captcha_text_image(img_path, img_name=None):
    """
    返回一个验证码的array形式和对应的字符串标签
    :return:tuple (str, numpy.array)
    """
    if img_name:
        # 标签
        label = img_name.split("-")[0]
    # 文件
    img_file = os.path.join(img_path, img_name)
    captcha_image = Image.open(img_file)
    return format_img_to_model(captcha_image), label


def format_folder(filepath, newFilepath=None):
    '''
    格式化文件夹图片
    :param filepath:
    :param newFilepath:
    :return:
    '''
    if not newFilepath:
        newFilepath = 'label/tmp'
    fileList = os.listdir(filepath)
    for i in fileList:
        im = Image.open(filepath + '/' + i)
        format_img(im, i, newFilepath)
    print('格式化完成')

def get_tag_img():
    """
    根据url获取合适的验证码
    :return:
    """
    r = requests.get(my_config.IMG_URL).content
    base64_str = base64.b64encode(r)
    base64_str = str(base64_str, 'utf8')
    # print(base64_str)
    im = img_utils.base64_to_image(base64_str)
    tag = model_precit.predict_img(im)
    return base64_str, tag

if __name__ == '__main__':
    format_folder('label/old', 'label/tmp')
    img_utils.split_train_and_test('label/tmp', my_config.train_path, my_config.test_path, my_config.char_set)


